package org.opens.aop;

import com.google.common.util.concurrent.RateLimiter;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.opens.util.ContextUtils;
import org.springframework.stereotype.Component;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

@Component
@Aspect
public class RateLimitAop {

    private RateLimiter rateLimiter = RateLimiter.create(1);

    @Pointcut("@annotation(org.opens.aop.aspect.RateLimitAspect)")
    public void serviceLimit() {

    }

    @Around("serviceLimit()")
    public Object around(ProceedingJoinPoint joinPoint) {
        Boolean flag = rateLimiter.tryAcquire();
        Object result = null;
        try {
            if (flag) {
                result = joinPoint.proceed();
            } else {
                HttpServletResponse response = ContextUtils.getHttpResponse();
                output(response, "访问太过频繁");
            }
        } catch (Throwable e) {
            e.printStackTrace();
        }
        return result;
    }

    public void output(HttpServletResponse response, String msg) throws IOException {
        ServletOutputStream outputStream = null;
        try {
            if (response != null) {
                response.setContentType("application/json;charset=UTF-8");
                outputStream = response.getOutputStream();
                outputStream.write(msg.getBytes("UTF-8"));
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (outputStream != null) {
                outputStream.flush();
            }
            if (outputStream != null) {
                outputStream.close();
            }
        }
    }

}
